#!/usr/bin/env python

import rospy,math,os
from std_msgs.msg import Int32 
# from geometry_msgs.msg import PoseStamped,Twist
from mavros_msgs.srv import CommandBool, SetMode, CommandVtolTransition, CommandLong, CommandVtolTransitionRequest, CommandHome
from mavros_msgs.msg import State, PositionTarget, ExtendedState
from geometry_msgs.msg import PoseStamped, TwistStamped,Vector3,Twist
import tf.transformations as tf

class ROS_INIT():
    def ros_init(self):
        rospy.init_node('uav_formation', anonymous=True)
        self.rate = rospy.Rate(20)   # Hz

        ############# ros 初始化 #############
        self.cmd_vel_pub = rospy.Publisher('/'+self.vehicle_type+str(self.uav_id)+'/cmd_vel', Twist, queue_size=10)

        rospy.Subscriber('/command', Int32, self.command_callback)
        angle_sub = [None]*self.uav_num
        pos_sub = [None]*self.uav_num
        vel_sub = [None]*self.uav_num
        for nm in range(self.uav_num):
            # if nm+1 == self.uav_id:
            #     continue

            # self.pose_sub[nm] = rospy.Subscriber('/ugv'+str(nm+1)+'/pose', PoseStamped, self.pose_callback,queue_size=1,callback_args=nm+1)
            angle_sub[nm] = rospy.Subscriber('/'+self.vehicle_type+str(nm+1)+'/angle', Vector3, self.uav_angle_callback, queue_size=1,callback_args=nm+1)
            pos_sub[nm] = rospy.Subscriber('/'+self.vehicle_type+str(nm+1)+'/pos', Vector3, self.uav_pos_callback, queue_size=1,callback_args=nm+1)
            vel_sub[nm] = rospy.Subscriber('/'+self.vehicle_type+str(nm+1)+'/vel', Twist, self.uav_vel_callback, queue_size=1,callback_args=nm+1)

    def uav_angle_callback(self, msg, nm):
        k = nm-1
        roll = msg.x /180 *math.pi                              # 滚转角 roll angle(rad) Front
        pitch = -msg.y / 180 * math.pi                          # 俯仰角 pitch angle(rad) Left (negative when rise)
        yaw = -(msg.z / 180 * math.pi- 0.5 * math.pi)          # 偏航角    
        yaw = math.atan2(math.sin(yaw), math.cos(yaw)) 
        self.uav_angle[k][0] = roll
        self.uav_angle[k][1] = pitch
        self.uav_angle[k][2] = yaw

    def uav_pos_callback(self, msg, nm):
        # print(msg)
        k = nm-1
        self.uav_pos[k][0] = msg.y
        self.uav_pos[k][1] = msg.x
        self.uav_pos[k][2] = -msg.z

    def uav_vel_callback(self,msg,nm):
        k = nm-1
        # print(msg)
        self.uav_vel[k][0] = msg.linear.x
        self.uav_vel[k][1] = msg.linear.y
        self.uav_vel[k][2] = msg.linear.z
        self.v = math.sqrt(msg.linear.x**2 + msg.linear.y **2+msg.linear.z**2)

    def send_cmd2(self, v, dpsi, h):
        self.roll_cmd = math.atan(-dpsi * self.v / self.g) * 180 / math.pi  # 偏航角率转滚转
        # self.roll_cmd = min(max(self.roll_cmd, -self.roll_thres), self.roll_thres)
        # 发布指令
        cmd_vel = Twist()
        cmd_vel.linear.x = v
        cmd_vel.linear.y = h
        cmd_vel.linear.z = self.roll_cmd     # 滚转
        cmd_vel.angular.y = 1     # 使能
        self.cmd_vel_pub.publish(cmd_vel)
    
    def command_callback(self, msg):
        self.command = msg.data
        print('Receive command: ', self.command)
        if self.command == 0:
            self.state = 0
            print('state: ', self.state, ' 初始化')
        elif self.command == 990:
            self.state = 2
            print('state: ', self.state, ' 停车')
        elif self.command == 50:
            self.state = 1
            print('state: ', self.state, ' 开始')
        else:
            print('state: ', self.state, ' 保持上一个状态')
class Formation(ROS_INIT):
    def __init__(self):
        ################# INIT #################################
        self.uav_id = rospy.get_param('uav_id', 1)    # 获取参数 -- 本机id
        print("本机是 {} 号飞机".format(self.uav_id))
        self.uav_num = rospy.get_param('uav_num', 1)    # 获取参数 -- 集结的无人机数量
        print("集结的无人机数量是 {} ".format(self.uav_num))
        self.g = 9.81
        self.roll_thres = 30
        self.uav_angle = [[0,0,0] for i in range(self.uav_num)]        
        self.uav_pos = [[0,0,0] for i in range(self.uav_num)]    
        self.uav_vel = [[0,0,0,0,0,0] for i in range(self.uav_num)]    
        self.tar_pos = [400,400]
        self.tar_r = 200
        self.vehicle_type = 'uav'
        self.jijie_over = [0 for i in range(self.uav_num)]
        self.v = 0                      # 和速度
        self.save_data = [[0,0,0,0] for i in range(self.uav_num)]       # 斜距，高低角，方位角，方位角变化率(方位角是东北天坐标系下，与正北的夹角)
        self.ref_point = [0,0,0]

        ################# ROS init ########################        
        self.command = 0     
        self.state = 0   
        self.ros_init()
        # rospy.Timer(rospy.Duration(1/20), self.state_timer)

    def main(self):
        start_time = 0.0
        v_sp = 25
        id  = self.uav_id-1
        folder = os.path.dirname(os.path.abspath(__file__))
        fname = folder + '/uav'+str(self.uav_id)+'的数据.txt'
        if False:
            print("ERROR!!!")
        else:
            while not rospy.is_shutdown():
                t = rospy.get_time() - start_time
                if self.state == 0:         # 初始化
                    start_time = rospy.get_time()
                    # print("等待")
                
                elif self.state == 1:       # 接收到集结点
                    dpsi_cmd = self.bili(self.v,self.uav_pos[id],self.uav_angle[id][2],self.tar_pos)
                    dx = self.uav_pos[id][0] - self.ref_point[0]
                    dy = self.uav_pos[id][1] - self.ref_point[1]
                    self.save_data[id][0] = self.distance3d(self.uav_pos[id],self.ref_point)
                    self.save_data[id][1] = math.atan2(self.uav_pos[id][2],self.distance(self.uav_pos[id],self.ref_point))
                    self.save_data[id][2] = math.atan2(dx,dy)
                    vx = self.v*math.cos(self.uav_angle[id][2])
                    vy = self.v*math.sin(self.uav_angle[id][2])
                    self.save_data[id][3] = (vx*dy - dx*vy)/(1+self.distance(self.uav_pos[id],self.ref_point)**2)
                    # 使用open()和write()保存字符串数据
                    self.save_simulation_data(
                                file_path=fname,
                                time=t,
                                data=self.save_data[id],
                                # velocity=(vx, vy, vz),
                                # acceleration=(ax, ay, az),
                            )
                    self.send_cmd2(v_sp,dpsi_cmd,100)


                self.rate.sleep()
    
    def bili(self,v,uav_pos,theta,tar_pos):
        '''
        v：无人机速度
        uav_pos,theta：无人机当前位置和姿态角
        tar_pos：目标位置
        '''
        N = 3
        r = self.distance(uav_pos,tar_pos)
        q = math.atan2(tar_pos[1]-uav_pos[1],tar_pos[0]-uav_pos[0])
        dq = -v/r*math.sin(theta-q)
        a = N*v*dq
        return a/v


    def state_timer(self,event):
        # if self.uav_id == 1:
            # print(self.distance(self.uav_pos[0],self.tar_pos)- self.tar_r-100)
            # print(self.distance(self.uav_pos[1],self.tar_pos)- self.tar_r-100)
            # print(self.distance(self.uav_pos[2],self.tar_pos)- self.tar_r-100)
            # print(self.distance(self.uav_pos[3],self.tar_pos)- self.tar_r-100)
            
        if abs(self.distance(self.uav_pos[0],self.tar_pos) - self.tar_r-100) <10:
            self.jijie_over[0] = 1
        if abs(self.distance(self.uav_pos[1],self.tar_pos) -self.tar_r-100) <10:
            self.jijie_over[1] = 1
        if abs(self.distance(self.uav_pos[2],self.tar_pos) -self.tar_r-100) <10:
            self.jijie_over[2] = 1
        if abs(self.distance(self.uav_pos[3],self.tar_pos) -self.tar_r-100) <10:
            self.jijie_over[3] = 1

    def loiter(self,v,pose,theta,center,r):
        # 逆时针跟圆 输出为弧度制
        # print(pose)
        uav_x = pose[0]
        uav_y = pose[1]
        uav_theta = theta
        center_x = center[0]
        center_y = center[1]
        l1 = 5 * v
        [x1,y1,x2,y2] = self.circle_intersect_circle(uav_x,uav_y,l1,center_x,center_y,r)
        if [x1,y1,x2,y2] == [None]*4:           #此处不考虑两点重合的情况
            # 圆心连线与盘旋圆交点，取离当前位置更近的点为跟踪点
            # print('无交点')
            k = (uav_y - center_y)/(uav_x - center_x)
            b0 = center_y - k * center_x
            [x1,y1,x2,y2,delta] = self.line_intersect_circle(uav_x, uav_y, l1,k,b0)
            dis_1 = self.distance(center,[x1,y1])
            dis_2 = self.distance(center,[x2,y2])
            if dis_1 < dis_2:
                inp_x = x1
                inp_y = y1
            else:
                inp_x = x2
                inp_y = y2
        else:           # 有交点取点逻辑写死了，没判断，无人机从绕圆左侧来，逆时针盘旋
            # vv = [self.v*math.cos(uav_theta),self.v*math.sin(uav_theta)]            # 速度向量
            # cc = [x1-uav_x,y1-uav_y]
            # if vv[0]*cc[0]+vv[1]*cc[1]>0:
            #     inp_x = x1
            #     inp_y = y1
            # else:
            #     inp_x = x2
            #     inp_y = y2
            inp_x = x2
            inp_y = y2
        # print('inp_xy: ',inp_x,inp_y)
        eta1 = math.atan2(inp_y - uav_y, inp_x - uav_x); # 位置角度
        # print('theta',uav_theta)
        eta2 = uav_theta                           # 速度角度
        eta = eta1 - eta2
        as_cmd = (2 * v * v / l1) * (math.sin(eta))
        w = as_cmd / v;	
        return w
    def circle_intersect_circle(self,x,y,R,a,b,S):
        # 求两圆相交交点
        d = self.distance([x,y],[a,b])
        # print(d,abs(R-S))
        if d > (R+S) or d <(abs(R-S)):
            return [None]*4
        elif d == 0:
            return [None]*4
        else:
            A = (R ** 2 - S ** 2 + d ** 2) / (2 * d)
            h = math.sqrt(R ** 2 - A ** 2)
            x2 = x + A * (a - x) / d
            y2 = y + A * (b - y) / d
            x3 = round(x2 - h * (b - y) / d, 2)
            y3 = round(y2 + h * (a - x) / d, 2)
            x4 = round(x2 + h * (b - y) / d, 2)
            y4 = round(y2 - h * (a - x) / d, 2)
            return [x3, y3, x4, y4]        

    def L1(self,v,pos,theta,exp_x,exp_y,exp_theta,bata,en):
        '''
        v：无人机巡航速度
        uav_x,uav_y,uav_theta：无人机当前位置和姿态角
        exp_x,exp_y,exp_theta：目标位置和期望入射角度
        bata：取值（0，pi/2），越大无人机越快收敛到目标点反向延长线
        en为True表示带入射角约束
        '''
        uav_x = pos[0]
        uav_y = pos[1]
        uav_theta = theta
        l1 = 5*v
        cel_x = uav_x
        cel_y = uav_y
        r = l1
        if en == True:
            k = math.tan(exp_theta)
            b0 = exp_y - k * exp_x
            # print(type(cel_y),type(cel_x))
            [x1,y1,x2,y2,delta] = self.line_intersect_circle(cel_x, cel_y, r,k,b0)
            if delta<0:
                # k = (exp_y - uav_y)/(exp_x - uav_x);
                if uav_y-k*uav_x > b0:               # 快速跟到直线
                    k = math.tan(exp_theta - bata)
                else:
                    k = math.tan(exp_theta + bata)
                b0 = cel_y - k * cel_x
                [x1,y1,x2,y2,delta] = self.line_intersect_circle(cel_x, cel_y, r,k,b0)
        else:
            k = (exp_y - cel_y)/(exp_x - cel_x)
            b0 = exp_y - k * exp_x
            [x1,y1,x2,y2,delta] = self.line_intersect_circle(cel_x, cel_y, r,k,b0)
        dis_1 = math.sqrt((exp_y - y1)*(exp_y - y1) + (exp_x - x1)*(exp_x - x1))
        dis_2 = math.sqrt((exp_y - y2)*(exp_y - y2) + (exp_x - x2)*(exp_x - x2))
        if dis_1 < dis_2:
            inp_x = x1
            inp_y = y1
        else:
            inp_x = x2
            inp_y = y2
        
        eta1 = math.atan2(inp_y - uav_y, inp_x - uav_x); # 位置角度
        eta2 = uav_theta                           # 速度角度
        eta = eta1 - eta2
        as_cmd = (2 * v * v / l1) * (math.sin(eta))
        w = as_cmd / v;	
        return w
    
    def line_intersect_circle(self,cel_x, cel_y, r,k,b0):
        # 直线与圆的方程化简为一元二次方程ax^2+bx+c=0,判断交点个数  
        # cel_x表示圆心的x坐标，cel_y表示圆心的y坐标
        a = k * k + 1
        # print(type(b0),type(cel_y),type(cel_x))
        b = 2 * k * (b0 - cel_y) - 2 * cel_x
        c = (b0 - cel_y) * (b0 - cel_y) + cel_x * cel_x - r * r
        # 判别式判断解
        delta = b * b - 4 * a * c
        if delta >= 0:
            x1 = (-b - math.sqrt(delta)) / (2 * a)
            y1 = k * x1 + b0
            x2 = (-b + math.sqrt(delta)) / (2 * a)
            y2 = k * x2 + b0
        else:
            x1 = 0
            y1 = 0
            x2 = 0
            y2 = 0
        return x1,y1,x2,y2,delta
    def distance(self,a,b):
        # print(a[0],b[0])
        x = a[0]-b[0]
        return math.sqrt((a[0]-b[0])**2+(a[1]-b[1])**2)
    def distance3d(self,a,b):
        # print(a[0],b[0])
        x = a[0]-b[0]
        return math.sqrt((a[0]-b[0])**2+(a[1]-b[1])**2 + (a[2]-b[2])**2)
    
    def save_simulation_data(self,file_path, time, data):
        """
        将单步仿真数据保存到txt文件

        参数:
            file_path: 保存文件的路径
            step: 当前仿真步
            data: 数据
        """
        # 如果文件不存在，创建文件并写入表头
        if not os.path.exists(file_path):
            with open(file_path, 'w') as f:
                headers = "Time,序号，斜距，高低角，方位角，方位角变化率"
                f.write(headers + "\n")
        
        # 写入当前步数据
        with open(file_path, 'a') as f:
            data_line = f"{time},{self.uav_id},{data[0]},{data[1]},{data[2]},{data[3]}"
            # data_line += f"{velocity[0]},{velocity[1]},{velocity[2]},"
            # data_line += f"{acceleration[0]},{acceleration[1]},{acceleration[2]}"
            f.write(data_line + "\n")
if __name__ == "__main__":
    formation = Formation()
    formation.main()